'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 3_study_doc_in_chain.py
* @Time: 2025/7/25
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()
# 创建提示词
mpt = ChatPromptTemplate.from_messages([
    ('system','你是一个聊天机器人, 根据用户的问题给出正确的答案,<context>{context}</context>'),
    ('human','{query}')
])
# 指定大模型
llm = ChatOpenAI(model= 'gpt-4o-mini', temperature=0)
# 创建链应用
chain = create_stuff_documents_chain(prompt= mpt, llm= llm)
# 模拟一个简单的文档列表
docs = [
    Document(page_content= '小明喜欢蓝色, 不喜欢红色'),
    Document(page_content= '小王喜欢黄色, 不喜欢蓝色'),
    Document(page_content= '小花喜欢白色, 不喜欢黑色色'),
]
# 构建调用链
content = chain.invoke({'query': '统计大家喜欢什么颜色', 'context': docs})
print(content)